import torch
import sys
import os
import numpy as np
import re

from torch.utils.data import Dataset

class MyDataset(Dataset):
    # 定义一个自己的数据集 用于训练遍历
    def __init__(self, root_path, type1="txt"):
        super().__init__()
        # 读取文件所有的列表
        self.root_path = root_path
        self.label_list  = os.listdir(root_path)[:2]  # 直选择 0 1 标签
        self.path = [ os.path.join(self.root_path, i) for i in self.label_list ]
        self.path_wav = [ os.path.join(root_path, "embed_audio", i) for i in self.label_list ]
        # print(self.path_wav)
        # 包括了所有文件的名字  使用txt气压数据  wav嵌入后的数据
        self.dir_list_txt = []          # 所有的txt文件名字 
        self.dir_list_txt_path = []     # 所有的txt路径
        for x in self.path:
            dir_list_txt = [i for i in os.listdir(x) if i.endswith(type1)]
            self.dir_list_txt += dir_list_txt
            self.dir_list_txt_path += [os.path.join(x, i) for i in dir_list_txt]
        
        self.dir_list_wav = []
        self.dir_list_wav_path = []
        for y in self.path_wav:
            dir_list_wav = [i for i in os.listdir(y)]
            self.dir_list_wav += dir_list_wav
            self.dir_list_wav_path += [os.path.join(y, i) for i in dir_list_wav]
        # print(self.dir_list_txt_path[:10])
    
    def __len__(self):
        return len(self.dir_list_wav)

    def __getitem__(self, index):
        # 返回数据集的切片
        # print(len(self.dir_list_txt), len(self.dir_list_txt_path))s
        file_txt = self.dir_list_txt_path[index]
        file_txt_label = file_txt.split("/")[-2]  # 找到倒数第二个位置 标签
        file_txt_name = self.dir_list_txt_path[index].split("/")[-1].split(".")[0]  # 获取文件名 但不包含后缀名字
        # print(file_txt, file_txt_name)
        pattern = re.compile(rf".*{re.escape(file_txt_label+'/'+file_txt_name)}.*")
        file_wav = [s for s in self.dir_list_wav_path if pattern.match(s)]
        # print(file_txt_name, file_wav_emebed)
        file_item_txt = np.loadtxt(file_txt)
        # print(file_wav)
        file_wav_emebed = np.load(file_wav[0])
        # print(file_txt_name)
        return (file_item_txt.reshape(1, 300), file_wav_emebed.reshape(1, 512)), np.array(int(file_txt_label))

if __name__ == "__main__":
    mydataset = MyDataset("./data/trainset")
    # print(len(mydataset))
    print(mydataset[10])
